import tkinter as tk
from PIL import Image, ImageOps, ImageDraw
import numpy as np
import tensorflow as tf
import io
import os

# 确保模型文件存在
if not os.path.exists('mnist_cnn_model.h5'):
    print("Error: Model file 'mnist_cnn_model.h5' not found!")
    exit()
else:
    model = tf.keras.models.load_model('mnist_cnn_model.h5')


class App:
    def __init__(self, master):
        self.master = master
        master.title("手写数字识别系统")
        master.resizable(False, False)
        # 顶部提示标签
        self.prompt = tk.Label(master, text="请在下方画板写数字", font=("Arial", 14))
        self.prompt.pack(pady=10)
        # 绘图画布
        self.canvas = tk.Canvas(master, width=350, height=280, bg='#F0F0F0')
        self.canvas.pack()
        # 底部控制面板
        ctrl_frame = tk.Frame(master)
        ctrl_frame.pack(pady=10)
        # 功能按钮
        self.recognize_btn = tk.Button(ctrl_frame, text="识别", command=self.predict)
        self.recognize_btn.grid(row=0, column=0, padx=5)
        self.clear_btn = tk.Button(ctrl_frame, text="清空", command=self.clear)
        self.clear_btn.grid(row=0, column=1, padx=5)
        # 结果显示
        self.result = tk.Label(master, text="识别结果：无", font=("Arial", 16), fg="blue")
        self.result.pack(pady=5)
        # 初始化绘图状态
        self.last_x, self.last_y = None, None
        self.canvas.bind('<ButtonPress-1>', self.on_press)
        self.canvas.bind('<B1-Motion>', self.on_move)
        self.canvas.bind('<ButtonRelease-1>', self.on_release)

    def on_press(self, event):
        """鼠标按下事件"""
        self.last_x, self.last_y = event.x, event.y
        self.current_path = []

    def on_move(self, event):
        """鼠标拖动事件"""
        if self.last_x and self.last_y:
            x, y = event.x, event.y
            self.canvas.create_line(self.last_x, self.last_y, x, y,
                                    width=6, fill='black', capstyle=tk.ROUND, smooth=True)
            self.last_x, self.last_y = x, y
            self.current_path.append((x, y))

    def on_release(self, event):
        """鼠标释放事件"""
        self.last_x, self.last_y = None, None

    def clear(self):
        """清空画布"""
        self.canvas.delete("all")
        self.result.config(text="识别结果：无")

    def predict(self):
        """执行识别"""
        try:
            # 创建临时画布（添加白色背景）
            image = Image.new("RGB", (280, 280), (255, 255, 255))
            draw = ImageDraw.Draw(image)

            # 重绘当前画布内容（优化绘制逻辑）
            for item in self.canvas.find_all():
                if self.canvas.type(item) == "line":
                    coords = self.canvas.coords(item)
                    # 收集所有点坐标
                    points = [(coords[i], coords[i + 1]) for i in range(0, len(coords), 2)]
                    if len(points) > 1:
                        # 使用原画布的真实绘制路径
                        draw.line(points, fill='black', width=12)

            # 关键修改1：颜色反转（匹配MNIST数据分布）
            image = image.convert("L")  # 转为灰度
            image = ImageOps.invert(image)  # 黑白反转（关键修正）

            # 关键修改2：智能裁剪（去除多余空白）
            bbox = image.getbbox()
            if bbox:  # 存在有效内容
                # 保持宽高比的裁剪
                cropped = image.crop(bbox)
                new_image = Image.new("L", (28, 28), 255)
                width, height = cropped.size
                scale = min(24 / width, 24 / height)  # 保留2像素边距
                resized = cropped.resize((int(width * scale), int(height * scale)), resample=Image.BILINEAR)
                paste_x = (28 - resized.width) // 2
                paste_y = (28 - resized.height) // 2
                new_image.paste(resized, (paste_x, paste_y))
                image = new_image
            else:  # 空白画布处理
                image = image.resize((28, 28))
            image = image.resize((28, 28), resample=Image.BILINEAR)  # 调整到模型输入尺寸
            image = np.array(image).flatten() / 255.0  # 展平并归一化
            input_data = np.expand_dims(image, axis=0)  # 添加batch维度 -> (1, 784)
            # 进行预测
            predictions = model.predict(input_data)
            predicted = np.argmax(predictions[0])
            confidence = predictions[0][predicted]
            self.result.config(text=f"识别结果：{predicted} ({confidence * 100:.1f}%)")
        except Exception as e:
            print("Error:", e)
            self.result.config(text="识别出错")

def main():
    root = tk.Tk()
    app = App(root)
    root.mainloop()


if __name__ == "__main__":
    main()